package edu.cmu.lti.lm.dt;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

public class Test {
	
	public static Set<Feature> features = Entry.getFeat();
	
	public static EntryCollection ent = null,test = null;
	
	public static void main(String[] argv) throws IOException{
		String corpus = argv[0];
		String testcorp = argv[1];
		features.addAll(initFeatures());
		ent = initEntries(corpus);
		test = initEntries(testcorp);
		Trainer tr = new Trainer(features, ent);
		TreeNode t = tr.doTrain(ent, 8,"T");
		Map<String,Double> dbl = EntryCollection.wordDistribute(ent); 
		t.printTree("");
		System.err.println("The Log likelihood of Tree         : " + logLikelihood(ent,t));
		System.err.println("The Log likelihood of Tree(Test)   : " + logLikelihood(test,t));
		System.err.println("The Log likelihood of Unigram      : " + logLikelihood2(ent,dbl));
		System.err.println("The Log likelihood of Unigram(Test): " + logLikelihood2(test,dbl));
	}
	
	private static double logLikelihood(EntryCollection cl, TreeNode tn){
		Iterator<Entry> it = cl.iterator();
		double d = 0;
		int k = 0;
		while(it.hasNext()){
			double s = tn.getLogProb(it.next());
			if(s > 0 ){
				continue;
			}else{
				d += s;
				k++;
			}
		}
			
		return d/k;
		
	}
	private static double l2 = Math.log(2);
	private static double logLikelihood2(EntryCollection cl,Map<String,Double> m){
		Iterator<Entry> it = cl.iterator();
		double d = 0;
		int k=0;
		while(it.hasNext())
		{
			Double ds = m.get(it.next().getWord());
			if(ds!=null){
				d += Math.log(ds)/l2;
				k++;
			}
		}
		return d/k;
		
	}

	private static EntryCollection initEntries(String file) throws IOException {
		EntryCollection ent = new EntryCollection();
		BufferedReader rd = new BufferedReader(new FileReader(file));
		String str;
		while((str = rd.readLine())!=null){
			String[] tokens = str.split("\\s+");
			History hist = new History();
			hist.add("<s>");
			for(int i = 0; i< tokens.length ; i++){
				Entry e = new Entry(tokens[i], hist);
				Iterator<Feature> it = features.iterator();
				while(it.hasNext()){
					it.next().addValue(hist);
				}
				ent.add(e);
				hist.add(0,tokens[i]);
			}
		}
		return ent;
	}

	/**
	 * init all features
	 * @return
	 */
	private static Set<Feature> initFeatures() {
		HashSet<Feature> fet = new HashSet<Feature>();
		fet.add(new BigramFeature());
		fet.add(new TrigramFeature());
		return fet;
	}
}
